package com.loveprogrammer.asm;

import com.alibaba.fastjson.JSON;
import org.objectweb.asm.*;
import org.objectweb.asm.commons.AdviceAdapter;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.lang.reflect.Method;

import static org.objectweb.asm.Opcodes.ASM5;

/**
 * @ClassName MethodTest
 * @Description 字节码增强第一站
 * @Author admin
 * @Date 2023/11/17 15:02
 * @Version 1.0
 */
public class MethodTest extends ClassLoader {

    // 测试的方法
    public Integer strToNumber(String str) {
        return Integer.parseInt(str);
    }

    public static void main(String[] args) throws Exception {
        byte[] bytes = new MethodTest().getBytes(MethodTest.class.getName());
        // 输出方法
        outputClazz(bytes, MethodTest.class.getSimpleName());

        // 测试方法
        Class<?> clazz = new MethodTest().defineClass("com.loveprogrammer.asm.MethodTest", bytes, 0, bytes.length);
        Method queryUserInfo = clazz.getMethod("strToNumber", String.class);

        // 正确入参；测试验证结果输出
        Object obj01 = queryUserInfo.invoke(clazz.newInstance(), "123");
        System.out.println("01 测试结果：" + obj01);

        // 异常入参；测试验证打印异常信息
        Object obj02 = queryUserInfo.invoke(clazz.newInstance(), "abc");
        System.out.println("02 测试结果：" + obj02);
    }

    /**
     * 字节码增强获取新的字节码
     * 以上这段代码就是 ASM 用于处理字节码增强的模版代码块。首先他会分别创建 ClassReader、ClassWriter，用于对类的加载和写入，这里的加载方式在构造方法中也提供的比较丰富。可以通过类名、字节码或者流的方式进行处理。
     *
     * 接下来是对方法的访问 MethodVisitor ，基本所有使用 ASM 技术的监控系统，都会在这里来实现字节码的注入。这里面目前用到了三个方法的，如下；
     *
     * onMethodEnter 方法进入时设置一些基本内容，比如当前纳秒用于后续监控方法的执行耗时。还有就是一些 Try 块的开始。
     * visitMaxs 这个是在方法结束前，用于添加 Catch 块。到这也就可以将整个方法进行包裹起来了。
     * onMethodExit 最后是这个方法退出时，用于 RETURN 之前，可以注入结尾的字节码加强，比如调用外部方法输出监控信息。
     * 基本上所有的 ASM 字节码增强操作，都离不开这三个方法。下面我就一步步来用指令将方法改造。
     */
    private byte[] getBytes(String className) throws IOException {

        ClassReader cr = new ClassReader(className);
        ClassWriter cw = new ClassWriter(cr, ClassWriter.COMPUTE_MAXS);

        cr.accept(new ClassVisitor(ASM5, cw) {

            @Override
            public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
                // 方法过滤
                if (!"strToNumber".equals(name))
                    return super.visitMethod(access, name, descriptor, signature, exceptions);

                MethodVisitor mv = super.visitMethod(access,name,descriptor,signature,exceptions);

                return new AdviceAdapter(ASM5,mv,access,name,descriptor) {
                    private Label from = new Label(),
                            to = new Label(),
                            target = new Label();

                    // 方法进入时修改字节码
                    @Override
                    protected void onMethodEnter() {
                        visitLabel(from);
                        visitTryCatchBlock(from,to,target,"java/lang/Exception");
                    }

                    // 访问局部变量和操作数栈
                    @Override
                    public void visitMaxs(int maxStack, int maxLocals) {
                        //标志：try块结束
                        mv.visitLabel(to);
                        //标志：catch块开始位置
                        mv.visitLabel(target);
                        mv.visitFrame(Opcodes.F_SAME1, 0, null, 1, new Object[]{"java/lang/Exception"});

                        // 异常信息保存到局部变量
                        int local = newLocal(Type.LONG_TYPE);
                        mv.visitVarInsn(ASTORE, local);

                        // 输出信息
                        mv.visitLdcInsn(className + "." + name);  // 类名.方法名
                        mv.visitVarInsn(ALOAD, local);
                        mv.visitMethodInsn(INVOKESTATIC, Type.getInternalName(MethodTest.class), "point", "(Ljava/lang/String;Ljava/lang/Throwable;)V", false);


                        // 抛出异常
                        mv.visitVarInsn(ALOAD, local);
                        mv.visitInsn(ATHROW);


                        super.visitMaxs(maxStack, maxLocals);
                    }

                    // 方法退出时修改字节码
                    @Override
                    protected void onMethodExit(int opcode) {
                        if((IRETURN <= opcode && opcode <= RETURN) || opcode == ATHROW) {
                            /***
                             * this.nextLocal，获取局部变量的索引值。这个值就让局部变量最后的值，也就是存放 ARETURN 的值(ARETURN，是返回对象类型，如果是返回 int 则需要使用 IRETURN)。
                             * ASTORE，将栈顶引用类型值保存到局部变量indexbyte中。这里就是把返回的结果，保存到局部变量。在你头脑中可以想象这有两块区域，一个是局部变量、一个是操作数栈。他们不断的进行压栈和操作。
                             * ALOAD，从局部变量indexbyte中装载引用类型值入栈。现在再将这个值放到操作数栈用，用于一会输出使用。
                             *
                             * 其实这里就是把返回值赋值给一个局部变量a 但是并不是直接给与，从操作流程的角度来说，是把返回值先赋值给一个局部变量，再把局部变量入栈
                             */
                            int nextLocal = this.nextLocal;
                            // 将栈顶引用类型值保存到局部变量indexbyte中
                            mv.visitVarInsn(ASTORE, nextLocal);
                            // 从局部变量indexbyte中装载引用类型值入栈
                            mv.visitVarInsn(ALOAD, nextLocal);

                            // 调用静态方法输出
                            mv.visitLdcInsn(className + "." + name);  // 类名.方法名
                            mv.visitVarInsn(ALOAD, nextLocal);
                            mv.visitMethodInsn(INVOKESTATIC, Type.getInternalName(MethodTest.class), "point", "(Ljava/lang/String;Ljava/lang/Object;)V", false);
                        }
                    }
                };
            }
        }, ClassReader.EXPAND_FRAMES);

        return cw.toByteArray();
    }
    /**
     * 输出字节码
     *
     * @param bytes     字节码
     * @param className 类名称
     */
    private static void outputClazz(byte[] bytes, String className) {
        // 输出类字节码
        FileOutputStream out = null;
        try {
            String pathName = MethodTest.class.getResource("/").getPath() + className + "SQM.class";
            out = new FileOutputStream(new File(pathName));
            System.out.println("ASM字节码增强后类输出路径：" + pathName + "\r\n");
            out.write(bytes);
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            if (null != out) try {
                out.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    // 在上面我们已经将返回内容赋值给参数，那么在 return 之前，我们就可以在添加一个方法来输出方法信息和出参了。
    public static void point(String methodName, Object response) {
        System.out.println("系统监控 :: [方法名称：" + methodName + " 输出信息：" + JSON.toJSONString(response) + "]\r\n");
    }

    public static void point(String methodName, Throwable throwable) {
        System.out.println("系统监控 :: [方法名称：" + methodName + " 异常信息：" + throwable.getMessage() + "]\r\n");
    }
}
